Table of Contents

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [1]:

    

from graphviz import Digraph
import uuid
from random import sample


class Node:
  
  def __init__(self, key):
    self.key = key
    self.left = None
    self.right = None
    self.parent = None
    self.is_root = False
    
  @property
  def depth(self):  
    # 该节点相对于根的深度, 必须为节点指定 parent, 只有根节点时算1
    return 1 if self.is_root else 1 + self.parent.depth
  
  @property
  def height(self):  
    # 该节点相对于自己子孙中最下一层的高度, 只有节点自身时算1
    height1 = self.left.height if self.left else 0
    height2 = self.right.height if self.right else 0
    return max(height1, height2) + 1

  @property
  def count(self):
    # 节点的所有子孙总数
    if self.left is None and self.right is None:
      return 1
    else:
      return self.left.count + self.right.count


  def __str__(self):
    return str(self.key)
  
  def __repr__(self):
    return "<{!r}> parent: {}, left: {}, right: {}".format(self.key, self.parent, self.left, self.right)

  

  
class Tree:
  
  def __init__(self):
    self.root = None

  def show(self, label=False, debug=False, size=None, node_surffix=lambda n: ''):
    if size is None:
      size = '{},{}'.format(self.root.height+1, self.root.height+1)
    self.dot = Digraph(comment='Binary Tree', graph_attr={'size': size})
    colors = ['aquamarine', 'khaki', 'skyblue', 'lightblue', 'orange', 'yellow', 'pink', 'lime', 'lightgreen']
    alpha = '22' if debug else '05'
    
    def pad_node(node_tag, child): 
      # 绘制跟child最大孩子数量的占位符, 
      # 否则自动排版后两个分支节点的位置会乱, 不好区分左右子树
      if child is None:
        return
      for i in range(2**(child.height-1)-1):
        pad_tag = str(uuid.uuid1())
        self.dot.node(pad_tag, "", color="#FF3366"+alpha) 
        self.dot.edge(node_tag, pad_tag, color="#999999"+alpha)

    def print_node(node, node_tag):  # 绘制以某个节点为根节点的二叉树
      if node.left is None and node.right is None:
        return
      left_tag = str(uuid.uuid1())  # 左节点的 uuid, 递归调用时用做连线 
      if node.left:
        pad_node(node_tag, node.left)
        self.dot.node(left_tag, str(node.left.key)+node_surffix(node.left), style='filled', color=color)  
        self.dot.edge(node_tag, left_tag, label=('L' if label else ''), color="#999999")   # 左节点与其父节点的连线
        pad_node(node_tag, node.left)
        print_node(node.left, left_tag)   # 递归调用, 打印下一层
      else:
        self.dot.node(left_tag, "", color="#FF3366"+alpha) # 占位的透明节点
        self.dot.edge(node_tag, left_tag, label=('L' if label else ''), color="#999999"+alpha)


      right_tag = str(uuid.uuid1())  
      if node.right:
        pad_node(node_tag, node.right)
        self.dot.node(right_tag, str(node.right.key)+node_surffix(node.right), style='filled', color=color)  
        self.dot.edge(node_tag, right_tag, label=('R' if label else ''), color="#999999")   # 左节点与其父节点的连线
        pad_node(node_tag, node.right)
        
        print_node(node.right, right_tag)   # 递归调用, 打印下一层
      else:
        self.dot.node(right_tag, "", color="#FF3366"+alpha) # 占位的透明节点
        self.dot.edge(node_tag, right_tag, label=('R' if label else ''), color="#999999"+alpha)
        
        
    if self.root is not None: 
      root_tag = str(uuid.uuid1())  # 根节点标签
      color = sample(colors,1)[0]    # 随机一个节点色
      self.dot.node(root_tag, str(self.root.key)+node_surffix(self.root), style='filled', color=color)  # 创建根节点
      print_node(self.root, root_tag)
    return self.dot


    

In [2]:

    

# 使用 Tree

tree = Tree()
tree.root = Node('A')
tree.root.left = Node('B')
tree.root.right = Node('C')
tree.root.left.right = Node('D')
tree.root.right.left = Node('E')
tree.root.right.right = Node('F')

tree.show(debug=False)


    

    
Out[2]:

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [3]:

    

class BinarySearchTree(Tree):
  def __init__(self):
    super().__init__()
    
  def insert(self, key):
    new_node = Node(key)
    if self.root is None:
      self.root = new_node
      return

    curr_node = self.root
    while curr_node is not None:
      parent_node = curr_node
      if new_node.key == curr_node.key:
        print('detect duplicate node "{}"! '.format(new_node))
        return
      elif new_node.key < curr_node.key:
        curr_node = curr_node.left
      else:
        curr_node = curr_node.right
    if new_node.key < parent_node.key:
      parent_node.left = new_node
    else:
      parent_node.right = new_node
    new_node.parent = parent_node
    return new_node
  
  def lookup(self, key):
    if self.root is None: return None
    curr_node = self.root
    while curr_node is not None:
      if curr_node.key == key: 
        return curr_node
      elif key < curr_node.key:
        curr_node = curr_node.left
      else:
        curr_node = curr_node.right
    return None


    

In [4]:

    

tree = BinarySearchTree()
tree.root = Node('C')
tree.insert('B')
tree.insert('A')
tree.insert('D')
tree.insert('F')
tree.insert('E')
tree.insert('A')  # 重复, 不能插入

tree.show(debug=False)


    

    




detect duplicate node "A"! 






    
Out[4]:

In [5]:

    

print(tree.lookup('C'))
print(tree.lookup('A'))
print(tree.lookup('F'))
print(tree.lookup('X'))


    

    




C
A
F
None

In [ ]:

    




    

In [6]:

    

class TraversalBinarySearchTree(BinarySearchTree):
  def __init__(self):
    super().__init__()
    
  def traversal(self, method='in_order'):
    if self.root is None: return []
    
    if method == 'in_order': return list(self.in_order(self.root))
    if method == 'pre_order': return list(self.pre_order(self.root))
    if method == 'post_order': return list(self.post_order(self.root))
    if method == 'dfs': return list(self.depth_first(self.root))
    if method == 'bfs': return list(self.breadth_first(self.root))
  
  
  def in_order(self, node):
    if not node: 
      return
    for n in self.in_order(node.left): yield n
    yield node
    for n in self.in_order(node.right): yield n

  def pre_order(self, node):
    if not node: 
      return
    yield node
    for n in self.pre_order(node.left): yield n
    for n in self.pre_order(node.right): yield n

  def post_order(self, node):
    if not node: 
      return
    for n in self.post_order(node.left): yield n
    for n in self.post_order(node.right): yield n
    yield node
      
  def depth_first(self, node):
    # 跟先序一样
    if not node: 
      return
    yield node
    for n in self.depth_first(node.left): yield n
    for n in self.depth_first(node.right): yield n
      
  def breadth_first(self, node):
    # 这个不会用递归, 只会用队列写
    queue = [node]
    while queue:
      current = queue.pop(0)
      if current.left: queue.append(current.left)
      if current.right: queue.append(current.right)
      yield current


    

In [7]:

    

# 构造以下这棵树, 演示遍历:
#          10
#        /    \
#       5      30
#     /       /  \
#    4       15   40
#  /
# 3
tree = TraversalBinarySearchTree()
tree.root = Node(10)
tree.insert(5)
tree.insert(30)
tree.insert(4)
tree.insert(3)
tree.insert(15)
tree.insert(40)
tree.show()


    

    
Out[7]:

In [8]:

    

print("中序:          ", [str(elem) for elem in tree.traversal('in_order')])
print("先序:          ", [str(elem) for elem in tree.traversal('pre_order')])
print("后序:          ", [str(elem) for elem in tree.traversal('post_order')])
print("深度优先 (先序):", [str(elem) for elem in tree.traversal('dfs')])
print("广度优先:       ", [str(elem) for elem in tree.traversal('bfs')])


    

    




中序:           ['3', '4', '5', '10', '15', '30', '40']
先序:           ['10', '5', '4', '3', '30', '15', '40']
后序:           ['3', '4', '5', '15', '40', '30', '10']
深度优先 (先序): ['10', '5', '4', '3', '30', '15', '40']
广度优先:        ['10', '5', '30', '4', '15', '40', '3']

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [9]:

    

class AVLTree(TraversalBinarySearchTree):
  def __init__(self):
    super().__init__()

  def get_balance_factor(self, node):
    h1 = node.left.height if node.left else 0
    h2 = node.right.height if node.right else 0
    return h1 - h2

  def bond(self, parent_node, node):
    if parent_node.key < node.key:
      parent_node.right = node
    else:
      parent_node.left = node
    node.parent = parent_node
  
  def left_rotate(self, node):
    father = node.parent  # node 如果有上游, 最后也需要修复链接关系
    son = node.right                     #   A                C  
    grandson = son.left                  #  / \              / \   
    son.left = node                      # B   C            A   Cr  
    node.parent = son                    #    / \    -->   / \   \  
    node.right = grandson                #   Cl  Cr       B  Cl  UC  
    if grandson: grandson.parent = node  #        \
    if father:                           #        UC
      self.bond(father, son)
    else:  # father 没有 parent, 所以  father 是 根节点
      self.root = son
      son.parent = None

  def right_rotate(self, node):
    father = node.parent  # node 如果有上游, 最后也需要修复链接关系 
    son = node.left                      #        A               B  
    grandson = son.right                 #       / \             / \   
    son.right = node                     #      B   C           Bl  A  
    node.parent = son                    #     / \      -->    /   / \  
    node.left = grandson                 #    Bl  Br          UB Br  C  
    if grandson: grandson.parent = node  #   /       
    if father:                           #  UB 
      self.bond(father, son)
    else:  # node 没有 parent, 所以 node 是 根节点
      self.root = son
      son.parent = None


    

In [10]:

    

tree = AVLTree()  # 构建 BST, 这个实现中, 插入节点时没考虑平衡
tree.insert('n')
tree.insert('i')
tree.insert('p')
tree.insert('g')
tree.insert('k')
tree.insert('f')
tree.insert('h')
node_surffix = lambda n: " ({:+1d})".format(tree.get_balance_factor(n))
tree.show(node_surffix=node_surffix)


    

    
Out[10]:

In [11]:

    

n = tree.lookup('n')
print(repr(n))
tree.right_rotate(n)  # 右旋 node_n, 这之后更平衡一些
tree.show(node_surffix=node_surffix)


    

    




<'n'> parent: None, left: i, right: p






    
Out[11]:

In [12]:

    

i = tree.lookup('i')
print(repr(i))
tree.right_rotate(i)   # 右旋 node_i, 只是演示右旋操作, 如果为了平衡这里不该这么旋转
tree.show(node_surffix=node_surffix)  # 这时 node_g 变成了 "问题节点"


    

    




<'i'> parent: None, left: g, right: n






    
Out[12]:

In [13]:

    

tree.left_rotate(tree.lookup('g'))    # 左旋 node_g
tree.show(node_surffix=node_surffix)


    

    
Out[13]:

In [14]:

    

tree = AVLTree()  # 构建另一个 BST, 演示 LR 型旋转
tree.insert('n')
tree.insert('i')
tree.insert('p')
tree.insert('g')
tree.insert('k')
tree.insert('j')   # 插入这个节点后, 导致 node_n 变为 +2, 之后需要使用 LR 修正
node_surffix = lambda n: " ({:+1d})".format(tree.get_balance_factor(n))
tree.show(node_surffix=node_surffix)


    

    
Out[14]:

In [15]:

    

tree.left_rotate(tree.lookup('i'))  # 首先左旋 node_i
tree.show(node_surffix=node_surffix)


    

    
Out[15]:

In [16]:

    

tree.right_rotate(tree.lookup('n'))  # 然后右旋 node_n
tree.show(node_surffix=node_surffix)


    

    
Out[16]:

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]: